#!/usr/bin/env python3

import subprocess
import re
import os
from textual.app import App, ComposeResult
from textual.widgets import Header, Footer, Button, Static, ListView, ListItem, Label, Input
from textual.containers import Container, Vertical
from textual.screen import Screen, ModalScreen
from textual import events

def strip_ansi_codes(text):
    """移除ANSI转义序列"""
    ansi_escape = re.compile(r'\x1b\[[0-9;]*m')
    return ansi_escape.sub('', text)

class PasswordScreen(ModalScreen):
    """密码输入屏幕"""
    
    def __init__(self, station, ssid):
        super().__init__()
        self.station = station
        self.ssid = ssid
    
    def compose(self) -> ComposeResult:
        yield Container(
            Label(f"输入网络 '{self.ssid}' 的密码:", id="password_label"),
            Input(placeholder="WiFi密码", password=True, id="password_input"),
            Button("连接", variant="primary", id="connect_btn"),
            Button("取消", id="cancel_btn"),
            id="password_dialog"
        )
    
    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "connect_btn":
            password = self.query_one("#password_input").value
            if password:
                self.dismiss((self.ssid, password))
        else:
            self.dismiss(None)

class NetworkScreen(Screen):
    """网络管理屏幕"""
    
    def __init__(self, station):
        super().__init__()
        self.station = station
        self.networks = []
    
    def compose(self) -> ComposeResult:
        yield Header()
        yield Container(
            Label(f"网卡 {self.station} - 选择网络:", id="network_question"),
            Button("🔄 扫描网络", id="scan_btn"),
            ListView(id="network_list"),
            id="network_container"
        )
        yield Footer()
    
    def on_mount(self) -> None:
        self.scan_networks()
    
    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "scan_btn":
            self.scan_networks()
    
    def scan_networks(self) -> None:
        """扫描WiFi网络"""
        try:
            # 执行扫描
            scan_result = subprocess.run(
                ["iwctl", "station", self.station, "scan"],
                capture_output=True,
                text=True,
                timeout=15
            )
            
            if scan_result.returncode != 0:
                self.show_error("扫描失败")
                return
            
            # 获取网络列表
            result = subprocess.run(
                ["iwctl", "station", self.station, "get-networks"],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            if result.returncode != 0:
                self.show_error("获取网络列表失败")
                return
            
            # 清理输出并解析
            clean_output = strip_ansi_codes(result.stdout)
            network_list = self.query_one("#network_list")
            network_list.clear()
            self.networks = []
            
            lines = clean_output.split('\n')
            for line in lines:
                line = line.strip()
                if not line or line.startswith('-') or 'SSID' in line:
                    continue
                
                parts = line.split()
                if len(parts) >= 2:
                    ssid = parts[0]
                    security = parts[1] if len(parts) > 1 else "open"
                    # 存储网络信息
                    self.networks.append({"ssid": ssid, "security": security})
                    # 显示在列表中
                    network_list.append(
                        ListItem(Label(f"{ssid} [{security}]"), id=f"network_{ssid}")
                    )
            
            if not self.networks:
                network_list.append(
                    ListItem(Label("未找到网络"), id="no_networks")
                )
                
        except subprocess.TimeoutExpired:
            self.show_error("扫描超时")
        except Exception as e:
            self.show_error(f"扫描异常: {str(e)}")
    
    def show_error(self, message):
        """显示错误信息"""
        network_list = self.query_one("#network_list")
        network_list.clear()
        network_list.append(
            ListItem(Label(message), id="error_message")
        )
    
    async def on_list_view_selected(self, event: ListView.Selected) -> None:
        """选择网络后的处理"""
        selected_item = event.item
        if selected_item and hasattr(selected_item, 'id') and selected_item.id:
            selected_id = selected_item.id
            if selected_id.startswith('network_'):
                ssid = selected_id.replace('network_', '')
                # 查找网络的安全类型
                network_info = next((n for n in self.networks if n["ssid"] == ssid), None)
                
                if network_info:
                    if network_info["security"] == "open":
                        # 开放网络，直接连接
                        self.connect_to_network(ssid, "")
                    else:
                        # 需要密码的网络，弹出密码输入框
                        password = await self.app.push_screen_wait(
                            PasswordScreen(self.station, ssid)
                        )
                        if password:
                            ssid, pwd = password
                            self.connect_to_network(ssid, pwd)
    
    def connect_to_network(self, ssid, password):
        """连接到指定网络"""
        try:
            if password:
                # 需要密码的连接
                cmd = ["iwctl", "station", self.station, "connect", ssid, "--passphrase", password]
            else:
                # 开放网络连接
                cmd = ["iwctl", "station", self.station, "connect", ssid]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=30
            )
            
            if result.returncode == 0:
                self.app.notify(f"已连接到 {ssid}", severity="information")
                self.app.pop_screen()  # 返回到主屏幕
            else:
                self.show_error(f"连接失败: {result.stderr}")
                
        except subprocess.TimeoutExpired:
            self.show_error("连接超时")
        except Exception as e:
            self.show_error(f"连接异常: {str(e)}")

class WiFiApp(App):
    """主应用"""
    
    BINDINGS = [("q", "quit", "退出"), ("r", "refresh", "刷新")]

    def compose(self) -> ComposeResult:
        yield Header()
        yield Container(Label("选择要操作的无线网卡:", id="question"))
        yield ListView(id="station_list")
        yield Footer()

    def on_mount(self) -> None:
        self.get_stations()

    def action_refresh(self) -> None:
        self.get_stations()

    def get_stations(self) -> None:
        """获取无线网卡列表"""
        try:
            result = subprocess.run(
                ["iwctl", "station", "list"],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            if result.returncode != 0:
                self.show_error(f"命令执行失败: {result.stderr}")
                return
            
            clean_output = strip_ansi_codes(result.stdout)
            station_list = self.query_one("#station_list")
            station_list.clear()
            
            lines = clean_output.split('\n')
            found_stations = False
            
            for line in lines:
                line = line.strip()
                if not line or line.startswith('-') or 'Name' in line or 'Devices' in line:
                    continue
                
                parts = line.split()
                if len(parts) >= 2:
                    station_name = parts[0]
                    station_state = parts[1]
                    unique_id = f"station_{station_name}"
                    station_list.append(
                        ListItem(Label(f"{station_name} [{station_state}]"), id=unique_id)
                    )
                    found_stations = True
            
            if not found_stations:
                station_list.append(
                    ListItem(Label("未找到无线网卡"), id="no_stations")
                )
        
        except subprocess.TimeoutExpired:
            self.show_error("命令执行超时")
        except Exception as e:
            self.show_error(f"异常: {str(e)}")

    def show_error(self, message):
        station_list = self.query_one("#station_list")
        station_list.clear()
        import time
        unique_id = f"error_{int(time.time() * 1000)}"
        station_list.append(
            ListItem(Label(message), id=unique_id)
        )

    def on_list_view_selected(self, event: ListView.Selected) -> None:
        """选择网卡后的处理"""
        selected_item = event.item
        if selected_item and hasattr(selected_item, 'id') and selected_item.id:
            selected_id = selected_item.id
            if selected_id.startswith('station_'):
                station_name = selected_id.replace('station_', '')
                # 推送到网络管理屏幕
                self.push_screen(NetworkScreen(station_name))

def main():
    if os.geteuid() != 0:
        print("请使用 sudo 运行: sudo python3 wifi-right.py")
        exit(1)
    
    try:
        subprocess.run(["systemctl", "is-active", "--quiet", "iwd"])
    except:
        print("请先启动 iwd 服务: sudo systemctl start iwd")
        exit(1)
    
    app = WiFiApp()
    app.run()

if __name__ == "__main__":
    main()
